# modules/ai_chat.py
import re

from glm import init_glm_ai, ask_glm
import streamlit as st
import matplotlib.pyplot as plt 


plt.rcParams['font.sans-serif'] = ['SimHei']  
plt.rcParams['axes.unicode_minus'] = False  

from glm import init_glm_ai, ask_glm 

def render_ai_chat(df):
    st.header("💬 智能招聘助手")

    with st.sidebar:
        api_key = st.text_input('智谱API Key:', type='password')
        client = init_glm_ai(api_key) if api_key else None

        # 新增：模型选择下拉菜单
        model_options = ['glm-4', 'glm-3-turbo', 'glm-4v'] # 您可以按需添加更多模型
        selected_model = st.selectbox("选择模型:", model_options, key='selected_model')

        # 新增：图表类型选择下拉菜单
        chart_type_options = ['饼图', '条形图'] # 可以根据工具函数支持的类型扩展
        selected_chart_type = st.selectbox("选择图表类型:", chart_type_options, key='selected_chart_type')
        
        if st.button("清除聊天记录"):
            st.session_state.messages = [{"role": "assistant", "content": "请问您想了解什么招聘信息？"}]
            st.rerun()

    if "messages" not in st.session_state:
        st.session_state.messages = [{"role": "assistant", "content": "请问您想了解什么招聘信息？"}]

    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.write(msg["content"])

    if prompt := st.chat_input("输入您的问题..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)

        with st.chat_message("assistant"):
            if not client:
                st.error("请先在侧边栏输入您的智谱API Key")
                return
            
            with st.spinner("思考中..."):
                messages_for_glm = st.session_state.messages
                
                # 调用 ask_glm 时传递所选模型和图表类型
                response_data = ask_glm(
                    client=client,
                    tools=load_tools(),
                    temperature=0.7,
                    messages=messages_for_glm, 
                    function_list=load_tools(), # 传递工具定义列表
                    dataset=df,
                    model_name=st.session_state.selected_model, # 使用会话状态中的模型
                    chart_type=st.session_state.selected_chart_type # 使用会话状态中的图表类型
                )
                
                if isinstance(response_data, dict) and "chart" in response_data:
                    st.write(response_data["text"])
                    if response_data["chart"]:
                        st.pyplot(response_data["chart"])
                        plt.close(response_data["chart"])
                    st.session_state.messages.append({"role": "assistant", "content": response_data["text"]})
                elif isinstance(response_data, str): # 处理纯文本回复或错误信息
                    st.write(response_data)
                    st.session_state.messages.append({"role": "assistant", "content": response_data})
                else: # 处理glm直接返回的消息对象 (如果未使用工具)
                    content_to_display = response_data.content if hasattr(response_data, 'content') else str(response_data)
                    st.write(content_to_display)
                    st.session_state.messages.append({"role": "assistant", "content": content_to_display})


def load_tools():
    """集成实验4的工具定义"""
    return [
        {
            "type": "function",
            "function": {
                "name": "get_average_salary",
                "description": "获取特定岗位的平均薪资信息",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "position_name": {"type": "string", "description": "岗位名称，如Java工程师"}
                    },
                    "required": ["position_name"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_city_positions",
                "description": "获取指定城市的岗位需求排名",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "city": {"type": "string", "description": "城市名称，如福州"}
                    },
                    "required": ["city"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_skill_distribution",
                "description": "获取指定技能关联的岗位分布及数量，并推荐相关技能",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "skill": {"type": "string", "description": "技能名称，如C++"}
                    },
                    "required": ["skill"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_city_specific_position",
                "description": "查询指定城市中特定岗位的需求量",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "city": {"type": "string", "description": "城市名称，如福州"},
                        "position_name": {"type": "string", "description": "岗位名称，如Java工程师"}
                    },
                    "required": ["city", "position_name"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_city_salary",
                "description": "获取指定城市中特定岗位的平均薪资",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "city": {"type": "string", "description": "城市名称，如福州"},
                        "position_name": {"type": "string", "description": "岗位名称，如Java工程师"}
                    },
                    "required": ["city", "position_name"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_position_skills",
                "description": "获取某岗位常见技能分布",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "position_name": {"type": "string", "description": "岗位名称，如嵌入式工程师"}
                    },
                    "required": ["position_name"]
                }
            }
        }
    ]


# glm.py
from zhipuai import ZhipuAI
import json

def init_glm_ai(api_key):
    return ZhipuAI(api_key=api_key)

def chatTools(client, messages, tools, model_name="glm-4", temperature=0.6):
    response = client.chat.completions.create(
        model=model_name,
        messages=messages,
        tools=tools,
        tool_choice="auto",
        temperature=temperature
    )
    return response.choices[0]

def chat_glmOnce(client, messages):
    response = client.chat.completions.create(
        model="glm-4",
        messages=messages # 传递完整的消息历史
    )
    return response.choices[0].message.content

def ask_glm(client, tools, temperature, messages, function_list, dataset, model_name="glm-4", chart_type="饼图"):
    # 调用智谱AI聊天工具，生成初步回复
    ack = chatTools(client, messages, tools, model_name=model_name, temperature=temperature)
    # 检查是否需要调用工具
    if ack.finish_reason == "tool_calls":
        # 提取工具调用信息
        tool_call = ack.message.tool_calls[0]
        function_name = tool_call.function.name# 工具函数名称
        function_args_dict = json.loads(tool_call.function.arguments)# 工具函数参数
        # 查找匹配的工具函数
        target_function = None
        for tool_def in function_list:
            if tool_def["function"]["name"] == function_name:
                try:
                    target_function = globals()[function_name]
                except KeyError:
                    return f"错误：在glm.py的全局作用域中未找到函数 {function_name}"
                break
        else:# 如果未找到匹配的工具定义，返回错误信息
            return "未找到匹配的工具定义"
        # 准备工具函数的参数
        final_args = {
            k: v for k, v in function_args_dict.items()
        }
        final_args["dataset"] = dataset# 添加数据集参数
        final_args["chart_type"] = chart_type# 添加图表类型参数
        # 根据工具函数名称调用相应的函数
        if function_name == "get_city_positions":
            from modules.ai_chat import get_city_positions
            result = get_city_positions(**final_args)
        elif function_name == "get_average_salary":
            from modules.ai_chat import get_average_salary
            result = get_average_salary(**final_args)
        elif function_name == "get_skill_distribution":
            from modules.ai_chat import get_skill_distribution
            result = get_skill_distribution(**final_args)
        elif function_name == "get_city_specific_position":
            from modules.ai_chat import get_city_specific_position
            result = get_city_specific_position(**final_args)
        elif function_name == "get_city_salary":
            from modules.ai_chat import get_city_salary
            result = get_city_salary(**final_args)
        elif function_name == "get_position_skills":
            from modules.ai_chat import get_position_skills
            result = get_position_skills(**final_args)
        else:
            return "未知的工具函数"
        return result
    return ack.message


# 工具函数定义
def get_city_positions(city, dataset, chart_type="饼图"):
    # 检查城市参数是否有效，若无效返回提示信息
    if not isinstance(city, str) or city.strip() == "":
        return {"text": "城市参数无效", "chart": None}
    # 过滤出指定城市的数据
    filtered_df = dataset[dataset['city'] == city]
    # 若筛选后的数据为空，返回提示信息
    if filtered_df.empty:
        return {"text": f"未找到{city}的数据", "chart": None}
    # 统计岗位出现频率，取前5个最常见的岗位
    counts = filtered_df['position'].value_counts().head(5)
    # 构造文本响应，包含需求量最大的岗位及前5岗位分布信息
    text_response = f"{city}需求量最大的岗位是：{counts.index[0]}。\n以下是需求量前5的岗位分布："
    # 创建图表
    fig, ax = plt.subplots()
    # 根据图表类型绘制对应的图表
    if chart_type == '饼图':
        ax.pie(counts, labels=counts.index, autopct='%1.1f%%', startangle=90)
        ax.axis('equal')  # 保证饼图为圆形
        plt.title(f'{city}岗位需求TOP5 (饼图)')
    elif chart_type == '条形图':
        counts.plot(kind='bar', ax=ax)
        ax.set_ylabel('数量')  # 设置y轴标签
        plt.title(f'{city}岗位需求TOP5 (条形图)')
        plt.xticks(rotation=45, ha='right')  # 设置x轴刻度标签旋转角度和对齐方式
        plt.tight_layout()  # 自动调整子图参数，使布局更合理
    else:
        # 若图表类型不支持，返回提示信息
        return {"text": text_response + "\n(不支持的图表类型)", "chart": None}
    # 返回文本响应和图表对象
    return {"text": text_response, "chart": fig}


def get_average_salary(position_name, dataset, chart_type="条形图"):
    # 过滤出指定岗位的数据
    filtered_df = dataset[dataset['position'] == position_name]
    # 若数据为空，返回提示信息
    if filtered_df.empty:
        return {"text": f"未找到{position_name}的薪资数据", "chart": None}
    # 计算平均薪资
    avg_salary = filtered_df['mean_salary'].mean()
    # 构造文本响应，包含该岗位的平均薪资
    text_response = f"{position_name}的平均薪资是：{avg_salary:.1f}k"
    # 创建图表
    fig, ax = plt.subplots()
    # 根据图表类型绘制对应的图表
    if chart_type == '条形图':
        ax.bar(['平均薪资'], [avg_salary], color='cornflowerblue')  # 设置柱状图颜色
        ax.set_ylabel('平均薪资 (k)')  # 设置y轴标签
        ax.set_title(f'{position_name}的平均薪资')  # 设置图表标题
    else:
        # 若图表类型不支持，返回提示信息
        return {"text": text_response, "chart": None}
    # 返回文本响应和图表对象
    return {"text": text_response, "chart": fig}


def get_skill_distribution(skill, dataset, chart_type="条形图"):
    # 对技能名称进行标准化处理（小写、去除特殊字符）
    normalized_skill = re.escape(skill.lower().strip())
    # 复制数据集以避免修改原始数据
    dataset = dataset.copy()
    # 将技能列表转换为小写，便于后续匹配
    dataset['clean_skills'] = dataset['skill_list'].fillna('').str.lower()
    # 过滤出包含指定技能的数据
    filtered_df = dataset[dataset['clean_skills'].str.contains(normalized_skill, regex=True, na=False)]
    # 若筛选后的数据为空，返回提示信息
    if filtered_df.empty:
        return {"text": f"未找到与{skill}相关的岗位数据", "chart": None}
    # 统计相关岗位的数量分布，取前5名
    position_counts = filtered_df['position'].value_counts().head(5)
    # 构造文本响应，包含相关岗位及其数量信息
    result_text = f"掌握{skill}可以从事以下岗位：\n"
    for pos, count in position_counts.items():
        result_text += f"{pos}: {count}个岗位\n"
    # 提取相关技能的推荐列表
    related_skills = filtered_df['skill_list'].dropna().str.split(',|，|/|\\|;|；').explode().str.strip()
    related_skills = related_skills[related_skills.str.lower() != skill.lower()]  # 排除指定技能本身
    top_related = related_skills.value_counts().head(5)
    if not top_related.empty:
        result_text += f"推荐你还可以学习：{', '.join(top_related.index)}"
    # 创建图表
    fig, ax = plt.subplots()
    # 根据图表类型绘制对应的图表
    if chart_type == '条形图':
        position_counts.plot(kind='bar', ax=ax)
        ax.set_title(f'与{skill}相关的岗位分布TOP5')  # 设置图表标题
        ax.set_ylabel('岗位数量')  # 设置y轴标签
        plt.xticks(rotation=45, ha='right')  # 设置x轴刻度标签旋转角度和对齐方式
        plt.tight_layout()  # 自动调整子图参数
    elif chart_type == '饼图':
        ax.pie(position_counts, labels=position_counts.index, autopct='%1.1f%%', startangle=90)
        ax.axis('equal')  # 保证饼图为圆形
        plt.title(f'与{skill}相关的岗位分布TOP5 (饼图)')  # 设置图表标题
    else:
        # 若图表类型不支持，返回提示信息
        return {"text": result_text + '\n(不支持的图表类型)', "chart": None}
    # 返回文本响应和图表对象
    return {"text": result_text, "chart": fig}


def get_city_specific_position(city, position_name, dataset, chart_type="条形图"):
    # 验证城市和岗位名称是否有效
    if not city or not position_name:
        return {"text": "需要同时指定城市和岗位", "chart": None}
    # 过滤出指定城市的数据
    city_df = dataset[dataset['city'] == city]
    # 若该城市的数据为空，返回提示信息
    if city_df.empty:
        return {"text": f"未找到{city}的任何数据", "chart": None}
    # 过滤出指定城市中该岗位的数据
    filtered = city_df[city_df['position'] == position_name]
    # 若该岗位的数据为空，返回提示信息
    if filtered.empty:
        return {"text": f"{city}暂无{position_name}岗位数据", "chart": None}
    # 计算该岗位在该城市所有岗位中的占比
    total_city_positions = len(city_df)
    count_specific_position = len(filtered)
    ratio = count_specific_position / total_city_positions * 100
    # 构造文本响应，包含该岗位的占比信息
    text_response = f"{city}的{position_name}需求占比：{ratio:.1f}%（共{count_specific_position}个岗位，该城市总岗位数{total_city_positions}）。"
    # 创建图表
    fig, ax = plt.subplots()
    # 定义标签和数据大小
    labels = [position_name, '其他岗位']
    sizes = [count_specific_position, total_city_positions - count_specific_position]
    if sizes[1] < 0: sizes[1] = 0  # 确保数据大小非负
    # 根据图表类型绘制对应的图表
    if chart_type == '条形图':
        ax.bar(labels, sizes, color=['skyblue', 'lightgray'])  # 设置柱状图颜色
        ax.set_ylabel('岗位数量')  # 设置y轴标签
        ax.set_title(f'{city}中{position_name}的需求占比')  # 设置图表标题
        plt.tight_layout()  # 自动调整子图参数
    elif chart_type == '饼图':
        ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=90)
        ax.axis('equal')  # 保证饼图为圆形
        plt.title(f'{city}中{position_name}的需求占比 (饼图)')  # 设置图表标题
    else:
        # 若图表类型不支持，返回提示信息
        return {"text": text_response + '\n(不支持的图表类型)', "chart": None}
    # 返回文本响应和图表对象
    return {"text": text_response, "chart": fig}


def get_city_salary(city, position_name, dataset, chart_type="条形图"):
    # 验证城市参数是否有效
    if not city or city not in dataset['city'].unique():
        return {"text": f"无效城市参数：{city}", "chart": None}
    # 验证岗位名称是否有效
    if not position_name:
        return {"text": "需要指定岗位名称", "chart": None}
    # 过滤出指定城市和岗位的数据
    filtered = dataset[(dataset['city'] == city) & (dataset['position'] == position_name)]
    # 若数据为空，返回提示信息
    if filtered.empty:
        return {"text": f"{city}暂无{position_name}薪资数据", "chart": None}
    # 计算该岗位在该城市的平均薪资
    avg_salary = filtered['mean_salary'].mean()
    # 构造文本响应，包含该岗位的平均薪资信息
    text_response = f"{city}的{position_name}的平均薪资是：{avg_salary:.1f}k"
    # 计算该城市所有岗位的平均薪资
    city_avg_salary = dataset[dataset['city'] == city]['mean_salary'].mean()
    # 创建图表
    fig, ax = plt.subplots()
    # 根据图表类型绘制对应的图表
    if chart_type == '条形图':
        salaries = [avg_salary, city_avg_salary]  # 定义薪资数据
        labels = [f'{position_name}平均薪资', f'{city}整体平均薪资']  # 定义标签
        ax.bar(labels, salaries, color=['cornflowerblue', 'lightcoral'])  # 设置柱状图颜色
        ax.set_ylabel('平均薪资 (k)')  # 设置y轴标签
        ax.set_title(f'{city} - {position_name}薪资对比')  # 设置图表标题
        plt.xticks(rotation=15, ha='right')  # 设置x轴刻度标签旋转角度和对齐方式
        plt.tight_layout()  # 自动调整子图参数
    else:
        # 若图表类型不支持，返回提示信息
        return {"text": text_response, "chart": None}
    # 返回文本响应和图表对象
    return {"text": text_response, "chart": fig}


def get_position_skills(position_name, dataset, chart_type="条形图"):
    # 过滤出包含指定岗位的数据
    filtered_df = dataset[dataset['position'].str.contains(position_name, case=False, na=False)]
    # 若数据为空，返回提示信息
    if filtered_df.empty:
        return {"text": f"未找到{position_name}相关的技能数据", "chart": None}
    # 拆分技能列表并统计出现频率，取前8个最常见的技能
    all_skills = filtered_df['skill_list'].dropna().str.split(',|，|/|\\|;|；').explode().str.strip()
    skill_counts = all_skills.value_counts().head(8)
    # 构造文本响应，包含该岗位的常见技能信息
    text_response = f"{position_name}岗位常见技能有：{', '.join(skill_counts.index)}"
    # 创建图表
    fig, ax = plt.subplots()
    # 根据图表类型绘制对应的图表
    if chart_type == '条形图':
        skill_counts.plot(kind='bar', ax=ax)
        ax.set_ylabel('出现次数')  # 设置y轴标签
        ax.set_title(f'{position_name}岗位常见技能TOP8')  # 设置图表标题
        plt.xticks(rotation=45, ha='right')  # 设置x轴刻度标签旋转角度和对齐方式
        plt.tight_layout()  # 自动调整子图参数
    elif chart_type == '饼图':
        ax.pie(skill_counts, labels=skill_counts.index, autopct='%1.1f%%', startangle=90)
        ax.axis('equal')  # 保证饼图为圆形
        plt.title(f'{position_name}岗位常见技能TOP8 (饼图)')  # 设置图表标题
    else:
        # 若图表类型不支持，返回提示信息
        return {"text": text_response + '\n(不支持的图表类型)', "chart": None}
    # 返回文本响应和图表对象
    return {"text": text_response, "chart": fig}